package vectorindex_v2

import (
	"context"
	"fmt"
	"strconv"

	"github.com/cloudwego/eino/schema"
	"github.com/milvus-io/milvus-sdk-go/v2/client"
	"github.com/milvus-io/milvus-sdk-go/v2/entity"
)

// MilvusRetrieverConfig 包含 Milvus 检索器的配置
type MilvusRetrieverConfig struct {
	// 与索引器共享的配置
	Address        string
	Username       string
	Password       string
	DatabaseName   string
	CollectionName string
	VectorDim      int
	MetricType     entity.MetricType

	// 检索特定配置
	DefaultTopK int
	DefaultExpr string
}

// WithMilvusRetrieverConfig 创建 Milvus 检索器配置选项
func WithMilvusRetrieverConfig(config *MilvusRetrieverConfig) RetrieveOption {
	return func(opt *RetrieveOptions) {
		// 这里简化处理，实际项目中可能需要更复杂的选项传递机制
	}
}

// MilvusRetriever 实现 VectorRetriever 接口
type MilvusRetriever struct {
	client client.Client
	config *MilvusRetrieverConfig
}

// NewMilvusRetriever 创建新的 Milvus 检索器
func NewMilvusRetriever(ctx context.Context, config *MilvusRetrieverConfig) (VectorRetriever, error) {
	// 验证必需配置
	if config.Address == "" {
		return nil, fmt.Errorf("milvus address is required")
	}
	if config.CollectionName == "" {
		return nil, fmt.Errorf("collection name is required")
	}

	// 设置默认值
	if config.DefaultTopK <= 0 {
		config.DefaultTopK = 10
	}
	if config.MetricType == 0 {
		config.MetricType = entity.L2
	}

	// 连接到 Milvus
	client, err := client.NewClient(ctx, client.Config{
		Address:      config.Address,
		Username:     config.Username,
		Password:     config.Password,
		DatabaseName: config.DatabaseName,
	})
	if err != nil {
		return nil, fmt.Errorf("failed to connect to milvus: %w", err)
	}

	// 确保集合存在
	exists, err := client.HasCollection(ctx, config.CollectionName)
	if err != nil {
		return nil, fmt.Errorf("failed to check collection existence: %w", err)
	}
	if !exists {
		return nil, fmt.Errorf("collection %s does not exist", config.CollectionName)
	}

	// 确保集合已加载
	err = client.LoadCollection(ctx, config.CollectionName, false)
	if err != nil {
		return nil, fmt.Errorf("failed to load collection: %w", err)
	}

	return &MilvusRetriever{
		client: client,
		config: config,
	}, nil
}

// Retrieve 实现 retriever.Retriever 接口的 Retrieve 方法
func (mr *MilvusRetriever) Retrieve(ctx context.Context, query string, opts ...RetrieveOption) ([]*schema.Document, error) {
	// 解析选项
	commonOpts := GetRetrieveCommonOptions(nil, opts...)

	// 确定 TopK
	topK := mr.config.DefaultTopK
	if commonOpts.TopK != nil && *commonOpts.TopK > 0 {
		topK = *commonOpts.TopK
	}

	// 确定过滤表达式
	expr := mr.config.DefaultExpr
	if commonOpts.Filter != nil && commonOpts.Filter["expr"] != nil {
		if exprStr, ok := commonOpts.Filter["expr"].(string); ok {
			expr = exprStr
		}
	}

	// 检查是否提供了嵌入器
	if commonOpts.Embedding == nil {
		return nil, fmt.Errorf("embedding component is required for retrieval")
	}

	// 生成查询向量
	vectors, err := commonOpts.Embedding.EmbedStrings(ctx, []string{query})
	if err != nil {
		return nil, fmt.Errorf("failed to embed query: %w", err)
	}
	if len(vectors) == 0 || len(vectors[0]) == 0 {
		return nil, fmt.Errorf("failed to generate valid embedding for query")
	}

	// 转换为 Milvus 向量格式
	queryVector := make([]float32, len(vectors[0]))
	for i, v := range vectors[0] {
		queryVector[i] = float32(v)
	}

	// 执行检索
	searchResult, err := mr.client.Search(
		ctx,
		mr.config.CollectionName,
		[]string{}, // 分区列表，空表示所有分区
		expr,
		[]string{"id", "content", "metadata"}, // 返回的字段
		entity.NewFloatVectorQuery([]float32(queryVector), mr.config.VectorDim),
		topK,
		entity.NewSearchParam("metric_type", string(mr.config.MetricType)),
	)
	if err != nil {
		return nil, fmt.Errorf("failed to search: %w", err)
	}

	// 解析结果
	var documents []*schema.Document

	for _, result := range searchResult { // 每个查询的结果（这里只有一个查询）
		for _, item := range result { // 每个返回的实体
			// 获取 ID
			idStr := strconv.FormatInt(item.ID, 10)

			// 获取内容和元数据
			var content string
			var metadataJSON string

			for i, fieldName := range result.Fields() {
				if i < len(item.Fields) {
					switch fieldName {
					case "content":
						if contentVal, ok := item.Fields[i].(string); ok {
							content = contentVal
						}
					case "metadata":
						if metadataVal, ok := item.Fields[i].(string); ok {
							metadataJSON = metadataVal
						}
					}
				}
			}

			// 创建文档
			doc := &schema.Document{
				ID:      idStr,
				Content: content,
			}

			// 解析元数据
			if metadataJSON != "" {
				if err := doc.UnmarshalMetaData(metadataJSON); err != nil {
					// 仅记录错误，不中断处理
					fmt.Printf("Warning: failed to unmarshal metadata: %v\n", err)
				}
			}

			// 添加相似度分数
			doc.MetaData["similarity_score"] = item.Score

			documents = append(documents, doc)
		}
	}

	return documents, nil
}

// Close 关闭 Milvus 连接
func (mr *MilvusRetriever) Close() error {
	if mr.client != nil {
		return mr.client.Close()
	}
	return nil
}
